#!/usr/bin/python
# -*- coding: utf-8 -*-
from functools import namedtuple
import time
import numpy as np

from env.env_5bits_2048 import Env_2048

env = Env_2048()


trans_table_entry = namedtuple("trans_table_entry", ["depth", "heuristic"])

class eval_state():
    def __init__(self) -> None:
        self.trans_table = dict()   # transposition table, to cache previously-seen moves, std::map
        self.maxdepth = 0
        self.curdepth = 0
        self.cachehits = 0
        self.moves_evaled = 0
        self.depth_limit = 0


SCORE_LOST_PENALTY = 200000.0
SCORE_MONOTONICITY_POWER = 4.0
SCORE_MONOTONICITY_WEIGHT = 47.0
SCORE_SUM_POWER = 3.5
SCORE_SUM_WEIGHT = 11.0
SCORE_MERGES_WEIGHT = 700.0
SCORE_EMPTY_WEIGHT = 270.0
heur_score_table = []        # 65536

def init_score_table():
    for row in range(2**16):
        line = [
            (row >>  0) & 0xf,
            (row >>  4) & 0xf,
            (row >>  8) & 0xf,
            (row >> 12) & 0xf]

        # Heuristic score
        sum = 0
        empty = 0
        merges = 0

        prev = 0
        counter = 0

        for rank in line:
            sum += pow(rank, SCORE_SUM_POWER)  # 行数据中各格数字的加权和

            if (rank == 0): 
                empty += 1
            else:
                if (prev == rank): 
                    counter += 1
                elif (counter > 0):
                    merges += 1 + counter
                    counter = 0
                prev = rank
 
        if (counter > 0): 
            merges += (1 + counter)      # merge 表示一行数据中可以合并的格数

        monotonicity_left = 0
        monotonicity_right = 0

        for i in range(1, 4): 
            if (line[i-1] > line[i]):
                monotonicity_left += pow(line[i-1], SCORE_MONOTONICITY_POWER) - pow(line[i], SCORE_MONOTONICITY_POWER)
            else:
                monotonicity_right += pow(line[i], SCORE_MONOTONICITY_POWER) - pow(line[i-1], SCORE_MONOTONICITY_POWER)


        heur_score_table.append( SCORE_LOST_PENALTY +       # 得分的基数
            SCORE_EMPTY_WEIGHT * empty +                    # 空格数 的加权分数
            SCORE_MERGES_WEIGHT * merges -                  # 可合并空格数 的加权分数
            SCORE_MONOTONICITY_WEIGHT * min(monotonicity_left, monotonicity_right) -      # 单调性越高分数越低
            SCORE_SUM_WEIGHT * sum                          # 空格数值的加权和，这里认为数值越大得分越低
            )                                               # 或许行数据越单调越不利于格的合并


def score_helper(board):
    return heur_score_table[int(np.sum(board[0]*np.array([2**12, 2**8, 2**4, 2**0])))] + \
           heur_score_table[int(np.sum(board[1]*np.array([2**12, 2**8, 2**4, 2**0])))] + \
           heur_score_table[int(np.sum(board[1]*np.array([2**12, 2**8, 2**4, 2**0])))] + \
           heur_score_table[int(np.sum(board[1]*np.array([2**12, 2**8, 2**4, 2**0])))]


# score a single board heuristically
def score_heur_board(board):
    return score_helper(board) + score_helper(board.T)


CPROB_THRESH_BASE = 0.0001
CACHE_DEPTH_LIMIT  = 15

# score over all possible tile choices and placements
def score_tilechoose_node(state, board, cprob):
    if (cprob < CPROB_THRESH_BASE or state.curdepth >= state.depth_limit): 
        state.maxdepth = max(state.curdepth, state.maxdepth)
        return score_heur_board(board)

    if (state.curdepth < CACHE_DEPTH_LIMIT):
        if(str(board) in state.trans_table):
            entry = state.trans_table[str(board)]
            if(entry.depth <= state.curdepth):
                state.cachehits += 1
                return entry.heuristic

    # 取平均值
    num_open = env.count_empty(board)
    cprob /= num_open

    a_index, b_index = np.where(board==0)
    n = len(a_index)
    res = 0
    for i in range(n):
        a = a_index[i]
        b = b_index[i]
        board_1 = np.copy(board)
        board_2 = np.copy(board)
        board_1[a, b] = 1
        board_2[a, b] = 2

        res += score_move_node(state, board_1, cprob * 0.9) * 0.9
        res += score_move_node(state, board_2, cprob * 0.1) * 0.1
    
    res = res / n

    if (state.curdepth < CACHE_DEPTH_LIMIT):
        entry = trans_table_entry(state.curdepth, res)
        state.trans_table[str(board)] = entry

    return res


def score_move_node(state, board, cprob):
    best = -999999999
    state.curdepth = state.curdepth+1      # 与环境交互，执行上下左右移动视作state深度加一
    for move in range(4):
        newboard = env.execute_move(move, board)
        state.moves_evaled = state.moves_evaled+1

        if (not np.all(board==newboard)):
            best = max(best, score_tilechoose_node(state, newboard, cprob))

    state.curdepth = state.curdepth-1

    return best


def score_toplevel_move(board, newboard, move):
    state = eval_state()
    state.depth_limit = max(3, env.count_distinct_tiles(board) - 2)

    a_time = time.time()
    res = score_tilechoose_node(state, newboard, 1.0) + 1e-6
    elapsed = time.time() - a_time

    if(flag):
        print("Move %d: result %f, eval'd %ld moves, (%d cache hits, %d cache size) in %.2f seconds, (maxdepth=%d)\n"%(move, res,
        state.moves_evaled, state.cachehits, len(state.trans_table), elapsed, state.maxdepth))

    return res


def find_best_move(board):
    best_score = -999999999
    best_move = -1

    for move in range(4): 
        newboard = env.execute_move(move, board)
        if(np.all(board == newboard)):
            continue

        res = score_toplevel_move(board, newboard, move)
        if(res > best_score):
            best_score = res
            best_move = move
    return best_move


def play_game_by_AI():
    # "penalty" for obtaining free 4 tiles
    board, penalty = env.init_board()         # 初始化 board
    
    init_score_table()    # 初始化行数据的启发式得分

    moveno = 0

    while(True):
        if(flag):
            print(board)
            print("Move Number #%d, Current scores: (heur %d, actual %d)\n"%(moveno, int(score_heur_board(board)), int(env.get_score(board, penalty))))

        # move为-1表示无可选择的移动, move为可移动选项
        move = find_best_move(board)
        if(move == -1):
            # "Illegal move!\n"
            break    # no legal moves

        newboard = env.execute_move(move, board)
        moveno = moveno+1

        board, _penalty = env.insert_tile(newboard)
        penalty += _penalty

    if(flag):
        print("Game over. Your score is %d. The highest rank you achieved was %d.\n"%(int(env.get_score(board, penalty)), int(max(board))))
    else:
        print(moveno, ",", env.get_score(board, penalty), ",", np.max(board))


if __name__ == '__main__':
    flag = True
    play_game_by_AI()
